"""
生产环境的Vanna AI数据查询服务 - 道麒项目
优化版本，包含配置管理、错误处理、日志记录等功能
"""

import asyncio
import dataclasses
import logging
import os
import time
from dataclasses import dataclass
from datetime import datetime
from functools import wraps
from logging.handlers import RotatingFileHandler, TimedRotatingFileHandler
from pathlib import Path
from typing import Any, Callable, Dict, Final, List, Optional, Tuple, TypeVar, Union, cast

import nest_asyncio
import uvicorn

# 模块级常量和配置
MAX_SQL_RESULT_ROWS = 500  # SQL查询返回的最大行数限制
MAX_SQL_LENGTH = 10000    # SQL语句最大长度限制
SQL_TIMEOUT_SECONDS = 30  # SQL执行超时时间

# 系统级业务上下文定义
SYSTEM_BUSINESS_CONTEXTS = [
    "数据查询系统基于Vanna框架构建，用于提供企业级数据分析功能",
    "系统支持多用户认证，包括Cookie和HTTP Header两种认证方式",
    "所有SQL查询必须经过安全检查，只允许SELECT类操作",
    "系统会为每个用户会话维护独立的上下文信息"
]

# 全局业务上下文列表，用于存储生成的业务表结构信息
business_contexts: List[str] = []

# 应用nest_asyncio以允许嵌套事件循环
nest_asyncio.apply()

# Vanna相关模块导入
from vanna import Agent
from vanna.capabilities.sql_runner import RunSqlToolArgs
from vanna.core.agent import AgentConfig
from vanna.core.enhancer import LlmContextEnhancer
from vanna.core.llm import LlmMessage
from vanna.core.registry import ToolRegistry
from vanna.core.tool import ToolContext, ToolResult
from vanna.core.user import UserResolver, User, RequestContext
from vanna.integrations.chromadb import ChromaAgentMemory
from vanna.integrations.mysql import MySQLRunner
from vanna.integrations.ollama import OllamaLlmService
from vanna.servers.fastapi import VannaFastAPIServer
from vanna.tools import RunSqlTool, VisualizeDataTool
from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool, SaveTextMemoryTool

# 配置日志
def setup_logging():
    """设置日志配置，包括控制台和文件输出"""
    # 创建logs目录（如果不存在）
    logs_dir = Path("logs")
    logs_dir.mkdir(exist_ok=True)
    
    # 获取当前日期作为日志文件名的一部分
    current_date = datetime.now().strftime("%Y%m%d")
    log_file = logs_dir / f"vanna_prod_{current_date}.log"
    
    # 配置日志格式
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建根日志记录器
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    
    # 清除现有的处理器
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)
    
    # 添加控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    root_logger.addHandler(console_handler)
    
    # 添加文件处理器（按大小轮转）
    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=10*1024*1024,  # 10MB
        backupCount=5,
        encoding='utf-8'
    )
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    root_logger.addHandler(file_handler)
    
    # 添加按时间轮转的处理器（每天午夜轮转）
    time_handler = TimedRotatingFileHandler(
        logs_dir / "vanna_prod_daily.log",
        when='midnight',
        interval=1,
        backupCount=30,
        encoding='utf-8'
    )
    time_handler.setLevel(logging.INFO)
    time_handler.setFormatter(formatter)
    root_logger.addHandler(time_handler)
    
    # 记录日志系统启动信息
    logger = logging.getLogger(__name__)
    logger.info(f"日志系统已初始化，日志文件: {log_file}")

# 初始化日志系统
setup_logging()
logger = logging.getLogger(__name__)

# 常用表配置 - 道麒项目核心业务表
COMMON_TABLES: Final[List[str]] = [
    'dr_crm_account',           # 客户账户表
    'dr_crm_potential',         # 销售机会表
    'dr_crm_sales_order',       # 销售订单表
    'dr_crm_sales_order_product', # 销售订单产品明细表
    'dr_crm_quotes',            # 报价单表
    'dr_crm_sale_contract'      # 销售合同表
]

# 业务上下文：用于存储动态生成的建表语句
business_contexts: List[str] = []

def handle_errors(func: Callable[..., Any]) -> Callable[..., Any]:
    """
    统一错误处理装饰器，自动捕获并记录函数执行过程中的异常
    
    捕获函数执行过程中的异常，记录详细的错误信息（包括堆栈跟踪），
    然后重新抛出带有用户友好错误消息的异常。保留原函数的元数据。
    
    Args:
        func: 被装饰的函数
        
    Returns:
        包装后的函数，保持原函数签名但添加了错误处理逻辑
        
    Notes:
        - 保留原始异常上下文，便于调试
        - 自动识别并处理同步/异步函数
        - 记录详细错误日志，包括方法名和错误信息
    """
    # 检查函数是否为异步函数
    if asyncio.iscoroutinefunction(func):
        @wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            """异步函数错误处理包装器"""
            try:
                return await func(*args, **kwargs)
            except Exception as e:
                # 记录详细错误信息，包括堆栈跟踪
                logger.error(
                    f"🔴 异步方法 {func.__name__} 执行失败: {str(e)}", 
                    exc_info=True
                )
                # 重新抛出异常，提供更友好的错误消息
                raise Exception(f"异步操作执行失败: {str(e)}") from e
        return async_wrapper
    else:
        @wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            """同步函数错误处理包装器"""
            try:
                return func(*args, **kwargs)
            except Exception as e:
                # 记录详细错误信息，包括堆栈跟踪
                logger.error(
                    f"🔴 同步方法 {func.__name__} 执行失败: {str(e)}", 
                    exc_info=True
                )
                # 重新抛出异常，提供更友好的错误消息
                raise Exception(f"同步操作执行失败: {str(e)}") from e
        return sync_wrapper


@dataclass
class DatabaseConfig:
    """
    数据库连接配置类，用于集中管理数据库连接参数
    
    存储数据库连接所需的全部参数，支持从环境变量加载配置。
    适用于MySQL数据库连接配置。
    
    Attributes:
        host: 数据库服务器主机地址，默认为内部服务器地址
        database: 数据库名称，用于指定连接的具体数据库
        user: 数据库用户名，用于身份验证
        password: 数据库密码，用于身份验证
        port: 数据库服务器端口号，MySQL默认3306
    """
    host: str
    database: str
    user: str
    password: str
    port: int
    
    @classmethod
    def from_env(cls) -> 'DatabaseConfig':
        """
        从环境变量加载数据库配置，支持灵活的部署配置
        
        Returns:
            DatabaseConfig: 配置完整的数据库配置对象
            
        Notes:
            环境变量优先级高于默认值，便于在不同环境中部署时灵活配置
        """
        return cls(
            host=os.getenv('DB_HOST', '192.168.1.238'),
            database=os.getenv('DB_NAME', 'hdp_manage_dq'),
            user=os.getenv('DB_USER', 'root'),
            password=os.getenv('DB_PASSWORD', 'Hc$msql%0226.'),
            port=int(os.getenv('DB_PORT', '3306'))
        )


@dataclass
class LLMConfig:
    """
    语言模型配置类
    
    存储LLM服务所需的配置参数，支持从环境变量加载配置。
    适用于不同的LLM服务提供商配置。
    
    Attributes:
        model: 语言模型名称
        host: LLM服务主机地址
    """
    model: str
    host: str
    
    @classmethod
    def from_env(cls) -> 'LLMConfig':
        """
        从环境变量加载LLM配置
        
        Returns:
            LLMConfig: 配置完整的LLM配置对象
        """
        return cls(
            model=os.getenv('LLM_MODEL', 'gpt-oss:120b'),
            host=os.getenv('LLM_HOST', 'http://zh.daoreach.com:11434')
        )


@dataclass
class ServerConfig:
    """
    服务器配置类
    
    存储应用服务器所需的配置参数，支持从环境变量加载配置。
    控制服务监听端口、代理执行迭代次数和LLM生成温度。
    
    Attributes:
        port: 服务器监听端口号
        max_tool_iterations: 工具最大迭代次数
        temperature: LLM生成温度参数
    """
    port: int
    max_tool_iterations: int
    temperature: float
    
    @classmethod
    def from_env(cls) -> 'ServerConfig':
        """
        从环境变量加载服务器配置
        
        Returns:
            ServerConfig: 配置完整的服务器配置对象
        """
        return cls(
            port=int(os.getenv('SERVER_PORT', '8001')),
            max_tool_iterations=int(os.getenv('MAX_TOOL_ITERATIONS', '30')),
            temperature=float(os.getenv('LLM_TEMPERATURE', '0.7'))
        )


class ProductionEditableRunSqlTool(RunSqlTool):
    """生产环境SQL执行工具，增强错误处理、安全检查和日志记录
    
    提供更严格的SQL安全验证、完善的错误处理机制和性能监控能力，
    确保SQL查询的安全执行和高效运行。
    """
    
    # 常量定义
    MAX_RESULT_ROWS = 1000
    
    # 允许执行的SQL命令类型（仅查询类操作）
    ALLOWED_SQL_COMMANDS = ["SELECT", "WITH", "EXPLAIN"]
    
    # 危险的SQL模式（修改数据或结构的操作）
    DANGEROUS_SQL_PATTERNS = [
        "DELETE", "UPDATE", "INSERT", "TRUNCATE", 
        "DROP", "ALTER", "CREATE", "REPLACE", 
        "GRANT", "REVOKE", "EXECUTE", "CALL"
    ]
    
    def _check_sql_safety(self, sql: str) -> bool:
        """
        验证SQL语句的安全性，防止执行危险操作
        
        Args:
            sql: 待执行的SQL语句
            
        Returns:
            bool: SQL语句是否安全可执行
        """
        if not sql or not sql.strip():
            logger.warning("空SQL语句被拒绝")
            return False
            
        # 标准化SQL为大写以进行不区分大小写的检查
        normalized_sql = sql.strip().upper()
        
        # 检查是否以允许的命令开头
        if not any(normalized_sql.startswith(cmd) for cmd in self.ALLOWED_SQL_COMMANDS):
            logger.warning(f"SQL语句不以允许的命令开头: {sql[:50]}...")
            return False
            
        # 检查是否包含危险模式
        if any(pattern in normalized_sql for pattern in self.DANGEROUS_SQL_PATTERNS):
            logger.warning(f"SQL语句包含危险模式: {sql[:50]}...")
            return False
            
        return True
    
    @handle_errors
    async def execute(self, context: ToolContext, args) -> ToolResult:
        """
        执行SQL查询，增加安全检查、错误处理和性能监控
        
        Args:
            context: 工具执行上下文
            args: SQL执行参数，包含待执行的SQL语句
            
        Returns:
            ToolResult: SQL执行结果
        """
        # 提取SQL语句
        sql = args.sql
        user_email = context.user.email
        
        logger.info(f"用户 {user_email} 请求执行SQL查询:\n{sql[:100]}..." if len(sql) > 100 else f"用户 {user_email} 请求执行SQL查询:\n{sql}")
        
        # 安全检查
        if not self._check_sql_safety(sql):
            safety_message = f"SQL查询被拒绝，因为它包含不安全的操作。只允许: {', '.join(self.ALLOWED_SQL_COMMANDS)}"
            logger.warning(f"用户 {user_email} 的SQL查询被安全机制拒绝")
            return ToolResult(success=False, value=safety_message)
        
        # 记录执行开始时间，用于性能监控
        start_time = time.time()
        
        try:
            # 执行SQL查询
            result = await super().execute(context, args)
            
            # 记录执行时间
            execution_time = time.time() - start_time
            logger.info(f"SQL查询执行完成，耗时: {execution_time:.2f}秒，用户: {user_email}")
            
            return result
        except Exception as e:
            # 记录详细错误信息
            execution_time = time.time() - start_time
            logger.error(f"SQL查询执行失败，耗时: {execution_time:.2f}秒，用户: {user_email}，错误: {str(e)}", exc_info=True)
            raise


class ProductionUserResolver(UserResolver):
    """
    生产环境用户解析器，增强安全性
    支持从多种来源获取用户身份信息，提供清晰的错误处理机制
    """
    
    # 默认用户信息常量，用于身份验证失败时的回退方案
    DEFAULT_EMAIL = "guest@example.com"  # 默认访客邮箱
    DEFAULT_GROUP = "user"               # 默认用户组权限
    ADMIN_EMAIL = "admin@example.com"    # 管理员邮箱
    ADMIN_GROUP = "admin"                # 管理员权限组
    
    async def resolve_user(self, request_context: RequestContext) -> User:
        """
        解析用户信息，实现多认证源优先级解析机制
        
        按优先级尝试不同的认证源:
        1. Cookie认证（优先级最高）
        2. Header认证（次之）
        3. 默认访客账户（保底方案）
        
        Args:
            request_context: 请求上下文对象，包含HTTP请求相关的认证信息
            
        Returns:
            User: 解析成功的用户对象，确保始终返回有效用户
            
        Notes:
            - 即使所有认证源都失败也会返回默认用户，不会抛出异常
            - 为每个认证尝试记录详细日志，便于问题排查
        """
        # 定义认证源方法列表，按优先级排序
        auth_sources = [
            self._get_user_from_cookie,
            self._get_user_from_header
        ]
        
        # 尝试每个认证源，直到找到有效的用户信息
        for auth_source in auth_sources:
            try:
                user_info = await auth_source(request_context)
                if user_info:
                    # 成功获取用户信息
                    email, group = user_info
                    logger.info(f"成功从{auth_source.__name__}获取用户: {email}, 组: {group}")
                    return User(id=email, email=email, group_memberships=[group])
            except Exception as e:
                logger.warning(f"从{auth_source.__name__}获取用户信息失败: {str(e)}")
        
        # 所有认证源都失败，使用默认访客账户作为保底方案
        logger.warning("未找到用户信息，使用默认访客账户")
        return self._create_default_user()
    
    async def _get_user_from_cookie(self, request_context: RequestContext) -> Union[Tuple[str, str], None]:
        """
        从Cookie中获取用户信息
        
        Returns:
            包含(email, group)的元组，如果未找到则返回None
        """
        user_email = request_context.get_cookie('vanna_email')
        if user_email:
            group = self.ADMIN_GROUP if user_email == self.ADMIN_EMAIL else self.DEFAULT_GROUP
            return (user_email, group)
        return None
    
    async def _get_user_from_header(self, request_context: RequestContext) -> Union[Tuple[str, str], None]:
        """
        从Header中获取用户信息
        
        Returns:
            包含(email, group)的元组，如果未找到则返回None
        """
        user_email = request_context.get_header('X-User-Email')
        if user_email:
            group = self.ADMIN_GROUP if user_email == self.ADMIN_EMAIL else self.DEFAULT_GROUP
            return (user_email, group)
        return None
    
    def _create_default_user(self) -> User:
        """
        创建默认访客用户
        
        Returns:
            默认的User对象
        """
        return User(
            id=self.DEFAULT_EMAIL,
            email=self.DEFAULT_EMAIL,
            group_memberships=[self.DEFAULT_GROUP]
        )


class ProductionLlmContextEnhancer(LlmContextEnhancer):
    """
    生产环境LLM上下文增强器，负责为系统提示词添加表结构信息、业务描述和查询指导。
    通过提供详细的数据库上下文，帮助LLM生成更准确、符合业务逻辑的SQL查询。
    """
    
    # 预定义业务表描述信息（作为类常量）
    TABLE_BUSINESS_DESCRIPTIONS = {
        'dr_crm_account': '客户账户表，存储客户的基本信息、联系方式和账户状态',
        'dr_crm_potential': '销售机会表，记录潜在客户和销售线索的详细信息和跟进状态',
        'dr_crm_sales_order': '销售订单表，包含订单的基本信息、客户信息、订单状态和金额',
        'dr_crm_sales_order_product': '销售订单产品明细表，记录订单中的具体产品信息、数量和单价',
        'dr_crm_quotes': '报价单表，存储给客户的报价信息、报价状态和有效期',
        'dr_crm_sale_contract': '销售合同表，记录与客户签订的正式合同信息、条款和履行状态'
    }
    
    def __init__(self, mysql_runner: MySQLRunner, agent_memory=None):
        """
        初始化LLM上下文增强器
        
        Args:
            mysql_runner: MySQL运行器实例，用于执行数据库查询
            agent_memory: 代理内存实例，用于存储和检索上下文信息（可选）
        """
        # 核心依赖注入
        self.mysql_runner = mysql_runner
        self.agent_memory = agent_memory
        
        # 缓存系统（私有属性表示内部状态）
        self._table_schema_cache = {}
        self._table_info_cache = {}
        self._all_tables = []
        
        # 初始化控制（使用下划线前缀表示私有状态）
        self._initialized = False
        self._init_lock = asyncio.Lock()
        
        # 为了向后兼容，保留实例属性引用
        self.table_business_descriptions = self.TABLE_BUSINESS_DESCRIPTIONS
    
    @handle_errors
    async def _ensure_initialized(self) -> None:
        """确保已初始化表结构信息，使用锁防止并发初始化
        
        采用双重检查锁定模式，确保在并发环境下只初始化一次。
        懒加载模式：仅在首次需要时初始化表信息，避免不必要的数据库查询。
        初始化包括：
        1. 获取数据库中所有可用表
        2. 预加载常用表结构
        3. 设置初始化标志
        4. 记录初始化日志
        
        Raises:
            Exception: 当初始化过程中遇到任何错误时
        """
        if not self._initialized:
            async with self._init_lock:
                if not self._initialized:
                    await self._get_all_tables()
                    await self._get_common_table_structures()
                    self._initialized = True
                    logger.info(f"表结构初始化完成，共加载 {len(self._table_schema_cache)} 个表结构")
    
    @handle_errors
    async def _get_all_tables(self) -> List[str]:
        """
        获取数据库中的所有表名
        
        只查询配置在COMMON_TABLES中的表，避免加载不必要的表结构。
        使用缓存机制避免重复查询，提高性能。
        
        Returns:
            List[str]: 数据库中符合条件的表名列表
            
        Raises:
            Exception: 当数据库查询失败时
        """
        if self._all_tables:
            return self._all_tables
            
        try:
            # 创建系统用户上下文，用于执行管理操作
            context = ToolContext(
                user=User(id="system", email="system@vanna.ai", group_memberships=['admin']),
                conversation_id="schema_loading",
                request_id="schema_request",
                agent_memory=self.agent_memory
            )
            
            # 构建查询语句，只获取配置的常用表
            tables_query = "SELECT table_name FROM information_schema.tables WHERE table_schema = '" + self.mysql_runner.database + "' AND table_name IN (" + ", ".join([f"'" + t + "'" for t in COMMON_TABLES]) + ")"
            result = await self.mysql_runner.run_sql(
                RunSqlToolArgs(sql=tables_query),
                context
            )
            
            if not result.empty:
                self._all_tables = result.iloc[:, 0].tolist()
                logger.info(f"获取到数据库表列表: {self._all_tables}")
            
        except Exception as e:
            logger.error(f"获取表列表时出错: {e}")
            raise
            
        return self._all_tables
    
    @handle_errors
    async def _get_common_table_structures(self) -> None:
        """
        获取常用表的结构
        
        预加载所有配置的常用表结构，提高后续查询响应速度
        """
        # 确保COMMON_TABLES中指定的表结构都已加载
        for table_name in COMMON_TABLES:
            # 如果表结构尚未缓存，则获取它
            if table_name not in self._table_schema_cache:
                await self._get_table_structure(table_name)
    
    @handle_errors
    async def _get_table_structure(self, table_name: str) -> str:
        """
        从数据库中获取表结构，包括表结构、注释和字段信息
        
        Args:
            table_name: 表名
            
        Returns:
            str: 表结构信息，包含完整的表定义和元数据
        """
        # 先检查缓存，避免重复查询
        if table_name in self._table_schema_cache:
            return self._table_schema_cache[table_name]
        
        # 创建系统用户上下文
        context = ToolContext(
            user=User(id="system", email="system@vanna.ai", group_memberships=['admin']),
            conversation_id="schema_loading",
            request_id="schema_request",
            agent_memory=self.agent_memory
        )
        
        # 获取表的完整信息（包括结构、注释、字段等）
        table_info = await self._get_table_info(table_name, context)
        
        if table_info:
            # 构建增强的表结构信息
            enhanced_schema = self._build_enhanced_table_schema(table_info)
            
            # 缓存结果，提高后续查询效率
            self._table_schema_cache[table_name] = enhanced_schema
            self._table_info_cache[table_name] = table_info
            
            logger.debug(f"获取表 {table_name} 增强结构成功")
            return enhanced_schema
        
        return ""
    
    @handle_errors
    async def _get_table_info(self, table_name: str, context: ToolContext) -> Dict[str, Any]:
        """
        获取表的完整信息，包括表结构、注释、字段信息等
        
        Args:
            table_name: 表名
            context: 工具上下文，包含用户信息和代理内存引用
        
        Returns:
            Dict[str, Any]: 包含表完整信息的字典，结构如下：
            {
                'name': str,              # 表名
                'database': str,          # 数据库名
                'create_sql': str,        # CREATE TABLE语句
                'table_comment': str,     # 表注释
                'business_description': str, # 业务描述
                'columns': List[Dict]     # 字段信息列表
            }
            
        Raises:
            Exception: 当数据库查询失败时
        """
        table_info = {
            'name': table_name,
            'database': self.mysql_runner.database,
            'create_sql': '',
            'table_comment': '',
            'business_description': self.table_business_descriptions.get(table_name, ''),
            'columns': []
        }
        
        # 1. 获取CREATE TABLE语句
        create_query = f"SHOW CREATE TABLE {table_name}"
        create_result = await self.mysql_runner.run_sql(
            RunSqlToolArgs(sql=create_query),
            context
        )
        if not create_result.empty and 'Create Table' in create_result.columns:
            table_info['create_sql'] = create_result['Create Table'].iloc[0]
        
        # 2. 获取表注释
        comment_query = f"""
        SELECT table_comment 
        FROM information_schema.tables 
        WHERE table_schema = '{self.mysql_runner.database}' 
        AND table_name = '{table_name}'
        """
        comment_result = await self.mysql_runner.run_sql(
            RunSqlToolArgs(sql=comment_query),
            context
        )
        if not comment_result.empty:
            table_info['table_comment'] = comment_result.iloc[0, 0] or ''
        
        # 字段信息已从create_table_sql中提取，无需额外查询
        pass
        
        return table_info
    
    @handle_errors
    def _build_enhanced_table_schema(self, table_info: dict) -> str:
        """
        构建增强的表结构描述，包含业务信息和字段说明
        
        Args:
            table_info: 表信息字典
            
        Returns:
            str: 格式化的增强表结构描述
        """
        table_name = table_info['name']
        database = table_info['database']
        
        # 使用列表收集各部分内容，最后join，提高性能和可读性
        schema_parts = [
            f"### 表名: {database}.{table_name}",
            "-" * 60
        ]
        
        # 添加业务描述
        if table_info['business_description']:
            schema_parts.extend([
                "",
                "## 业务描述",
                table_info['business_description']
            ])
        
        # 添加表注释
        if table_info['table_comment']:
            schema_parts.extend([
                "",
                "## 表注释",
                table_info['table_comment']
            ])
        
        # 添加字段信息表格（如果有字段数据）
        if table_info['columns']:
            schema_parts.extend([
                "",
                "## 字段信息",
                f"| {'字段名':<20} | {'数据类型':<20} | {'字段注释':<40} |",
                f"|{'-'*22}|{'-'*22}|{'-'*42}|"
            ])
            
            for column in table_info['columns']:
                schema_parts.append(
                    f"| {column['name']:<20} | {column['type']:<20} | {column['comment'][:37]:<40} |"
                )
        
        # 添加使用说明
        schema_parts.extend([
            "",
            "## 使用说明",
            f"- 表名格式: {database}.{table_name}",
            "- 请严格使用表中定义的字段名",
            "- 根据业务需求可能需要与其他表进行JOIN操作"
        ])

        # 日志记录 schema_parts 数据
        logger.info(f"schema_parts 数据: {schema_parts}")
        
        return "\n".join(schema_parts)

    @handle_errors
    async def generate_business_contexts(self, table_names: Optional[list[str]] = None) -> list[str]:
        """
        动态生成业务表的建表语句上下文，并返回列表

        说明：
        - 优先使用传入的 `table_names`；未指定时使用 COMMON_TABLES；若为空则回退到全部表
        - 每条上下文包含完整的 `CREATE TABLE` 语句（字段、类型、约束等）
        - 用于初始化代理内存，提供结构化的业务表信息
        """
        await self._ensure_initialized()

        names: list[str] = []
        if table_names:
            names = table_names
        elif COMMON_TABLES:
            names = COMMON_TABLES
        else:
            names = self._all_tables or []

        contexts: list[str] = []
        for name in names:
            schema = await self._get_table_structure(name)
            if schema:
                contexts.append(schema)

        # 若未能获取到指定表的结构，回退尝试数据库中的其他表（最多3个）
        if not contexts and self._all_tables:
            for name in self._all_tables:
                schema = await self._get_table_structure(name)
                if schema:
                    contexts.append(schema)
                if len(contexts) >= 3:
                    break

        return contexts
    

    @handle_errors
    async def enhance_system_prompt(
        self,
        system_prompt: str,
        user_message: str,
        user: User
    ) -> str:
        """
        增强系统提示词，添加表结构信息和查询指导
        
        Args:
            system_prompt: 原始系统提示词
            user_message: 用户消息
            user: 用户对象
            
        Returns:
            str: 增强后的系统提示词
        """
        # 确保表信息已初始化
        await self._ensure_initialized()
        
        database_name = self.mysql_runner.database
        
        # 构建增强提示词，采用更清晰的结构
        enhanced_prompt = [
            "# 数据查询助手 - 系统配置",
            "",
            "## 数据库基本信息",
            f"- 数据库类型: MySQL",
            f"- 数据库名: {database_name}",
            "",
            "## 核心业务表详细信息"
        ]
        
        # 添加核心业务表信息
        table_count = 0
        for table_name in COMMON_TABLES:
            # 确保表结构已加载
            if table_name not in self._table_schema_cache:
                await self._get_table_structure(table_name)
            
            if table_name in self._table_schema_cache:
                # 获取表的业务描述
                business_desc = self.table_business_descriptions.get(table_name, '')
                
                # 构建表信息部分
                table_info = [
                    f"### 表名: {table_name}",
                    f"**业务描述**: {business_desc}",
                    "**使用说明**:",
                    f"- 表名格式: {database_name}.{table_name}",
                    "- 请严格使用表中定义的字段名",
                    "- 根据业务需求可能需要与其他表进行JOIN操作",
                    ""
                ]
                
                enhanced_prompt.extend(table_info)
                table_count += 1
        
        # 添加SQL查询指导和用户信息
        enhanced_prompt.extend([
            "## SQL查询指导",
            "1. 请严格按照SQL语法规则生成有效的MySQL查询语句",
            "2. 查询时请使用完整的表名，格式为 `数据库名.表名`",
            "3. 字段名请使用原始名称，不要添加引号或别名，除非必要",
            "4. 请确保生成的SQL语句可以直接在数据库中执行，不包含占位符",
            "5. 对于复杂查询，请考虑性能优化，避免全表扫描",
            ""
        ])
        
        # 添加数量查询特殊指导（仅当用户询问数量相关问题时）
        if any(keyword in user_message.lower() for keyword in ['多少', '数量', '总数', 'count']):
            enhanced_prompt.extend([
                "## 数量查询特殊指导",
                "- 当用户询问数量相关问题时，请生成简单的COUNT查询",
                f"- 正确示例: SELECT COUNT(*) AS total_count FROM {database_name}.表名",
                "- 不要添加额外的WHERE条件，除非用户明确指定了过滤条件",
                "- 为结果字段提供有意义的别名，如total_count、customer_count等",
                ""
            ])
        
        # 添加用户信息和原始系统提示词
        enhanced_prompt.extend([
            "## 用户信息",
            f"- 当前用户: {user.email}",
            f"- 用户组: {', '.join(user.group_memberships)}",
            "",
            "## 原始系统提示词",
            system_prompt
        ])
        
        # 构建最终提示词
        final_prompt = "\n".join(enhanced_prompt)
        
        # 记录日志
        logger.info(f"已增强系统提示词，添加了{table_count}个核心表的信息")
        logger.info(f"用户消息: {user_message}")
        
        return final_prompt

    @handle_errors
    async def enhance_user_messages(
        self,
        messages: list[LlmMessage],
        user: User
    ) -> list[LlmMessage]:
        """
        增强用户消息，添加用户上下文信息
        
        Args:
            messages: 原始消息列表
            user: 用户对象
            
        Returns:
            list[LlmMessage]: 增强后的消息列表，包含用户身份信息
        """
        messages.append(LlmMessage(
            role="system", 
            content=f"[系统上下文信息：当前登录用户为 {user.email}，所属组为 {', '.join(user.group_memberships)}。此信息仅用于身份验证和权限控制，不是用户查询条件。]"
        ))
        return messages


@handle_errors
async def create_production_agent() -> Agent:
    """
    创建并配置生产环境的Vanna代理实例
    
    Returns:
        Agent: 完全配置好的Vanna代理实例，包含数据库连接、工具注册和上下文增强
    """
    logger.info("开始创建生产环境Vanna代理...")
    
    # 1. 加载配置
    configs = _load_all_configurations()
    logger.info("配置加载完成")
    
    # 2. 初始化核心组件
    llm_service = _initialize_llm_service(configs['llm'])
    db_tool = _initialize_database_tool(configs['db'])
    agent_memory = _initialize_agent_memory()
    user_resolver = ProductionUserResolver()
    
    # 3. 注册工具
    tool_registry = _register_required_tools(db_tool)
    
    # 4. 配置并创建代理
    agent = _create_and_configure_agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        user_resolver=user_resolver,
        agent_memory=agent_memory,
        db_tool=db_tool,
        server_config=configs['server']
    )
    
    # 5. 预加载表结构信息
    await _preload_table_structures(agent)
    
    logger.info("生产环境Vanna代理创建完成")
    return agent


def _load_all_configurations() -> Dict[str, Any]:
    """
    加载所有必要的配置信息
    
    Returns:
        Dict[str, Any]: 包含各类配置的字典
    """
    logger.info("加载代理所需的所有配置...")
    return {
        'db': DatabaseConfig.from_env(),
        'llm': LLMConfig.from_env(),
        'server': ServerConfig.from_env()
    }


def _initialize_llm_service(config: LLMConfig) -> OllamaLlmService:
    """
    初始化LLM服务实例
    
    Args:
        config: LLM配置对象
        
    Returns:
        OllamaLlmService: 配置好的LLM服务实例
    """
    logger.info(f"初始化LLM服务: {config.model} @ {config.host}")
    return OllamaLlmService(
        model=config.model,
        host=config.host
    )


def _initialize_database_tool(config: DatabaseConfig) -> ProductionEditableRunSqlTool:
    """
    初始化数据库工具和连接
    
    Args:
        config: 数据库配置对象
        
    Returns:
        ProductionEditableRunSqlTool: 配置好的SQL执行工具
    """
    logger.info(f"初始化数据库连接: {config.user}@{config.host}:{config.port}/{config.database}")
    
    return ProductionEditableRunSqlTool(
        sql_runner=MySQLRunner(
            host=config.host,
            database=config.database,
            user=config.user,
            password=config.password,
            port=config.port
        )
    )


def _initialize_agent_memory() -> ChromaAgentMemory:
    """
    初始化代理内存实例
    
    Returns:
        ChromaAgentMemory: 配置好的内存实例
    """
    logger.info("初始化代理内存...")
    
    # 使用Path确保跨平台兼容性
    persist_dir = str(Path("./chroma_db"))
    
    return ChromaAgentMemory(
        collection_name="vanna_memory",
        persist_directory=persist_dir
    )


def _register_required_tools(db_tool: ProductionEditableRunSqlTool) -> ToolRegistry:
    """
    注册所有必要的工具到工具注册表
    
    Args:
        db_tool: 数据库工具实例
        
    Returns:
        ToolRegistry: 注册了所有必要工具的工具注册表
    """
    logger.info("注册代理工具...")
    
    tools = ToolRegistry()
    tools.register_local_tool(db_tool, access_groups=['admin', 'user'])
    tools.register_local_tool(SaveQuestionToolArgsTool(), access_groups=['admin'])
    tools.register_local_tool(SearchSavedCorrectToolUsesTool(), access_groups=['admin', 'user'])
    tools.register_local_tool(SaveTextMemoryTool(), access_groups=['admin', 'user'])
    tools.register_local_tool(VisualizeDataTool(), access_groups=['admin', 'user'])
    
    logger.info("成功注册 5 个工具")
    return tools


def _create_and_configure_agent(
    llm_service: OllamaLlmService,
    tool_registry: ToolRegistry,
    user_resolver: UserResolver,
    agent_memory: ChromaAgentMemory,
    db_tool: ProductionEditableRunSqlTool,
    server_config: ServerConfig
) -> Agent:
    """
    配置并创建代理实例
    
    Args:
        llm_service: LLM服务实例
        tool_registry: 工具注册表
        user_resolver: 用户解析器
        agent_memory: 代理内存
        db_tool: 数据库工具
        server_config: 服务器配置
        
    Returns:
        Agent: 配置完成的代理实例
    """
    logger.info("配置代理参数...")
    
    # 创建代理配置
    agent_config = AgentConfig(
        max_tool_iterations=server_config.max_tool_iterations,
        stream_mode=True,
        temperature=server_config.temperature,
        include_tool_names=True,
        auto_save_context=True,
        max_tokens=None
    )
    
    # 创建上下文增强器
    context_enhancer = ProductionLlmContextEnhancer(
        mysql_runner=db_tool.sql_runner, 
        agent_memory=agent_memory
    )
    
    # 创建代理
    logger.info("创建代理实例...")
    return Agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        user_resolver=user_resolver,
        agent_memory=agent_memory,
        llm_context_enhancer=context_enhancer,
        config=agent_config
    )


async def _preload_table_structures(agent: Agent) -> None:
    """
    预加载业务表结构信息到代理内存
    
    Args:
        agent: 代理实例
    """
    logger.info("预加载业务表结构信息...")
    
    try:
        # 获取上下文增强器
        enhancer: ProductionLlmContextEnhancer = agent.llm_context_enhancer  # type: ignore
        
        # 生成业务表上下文
        contexts = await enhancer.generate_business_contexts(COMMON_TABLES)
        
        # 保存到全局变量供其他测试/脚本使用
        global business_contexts
        business_contexts = contexts
        
        # 创建管理员上下文以写入内存
        admin_user = User(id="admin", email="admin@example.com", group_memberships=['admin'])
        ctx = ToolContext(
            user=admin_user,
            conversation_id="business_schema_bootstrap",
            request_id="bootstrap_request",
            agent_memory=agent.agent_memory
        )
        
        # 将表结构写入内存
        cached_count = 0
        for text in contexts:
            await agent.agent_memory.save_text_memory(content=text, context=ctx)
            cached_count += 1
            logger.debug(f"已缓存表结构: {text[:100]}...")
        
        logger.info(f"成功预加载 {cached_count} 个业务表结构")
        
    except Exception as e:
        logger.error(f"预加载表结构失败: {str(e)}", exc_info=True)
        logger.warning("代理将继续运行，但可能影响SQL生成准确性")


async def create_app():
    """
    创建应用实例，用于uvicorn启动
    
    Returns:
        FastAPI应用实例
    """
    # 创建代理
    agent = await create_production_agent()
    
    # 创建并返回FastAPI应用
    server = VannaFastAPIServer(agent=agent)
    app = server.create_app()
    logger.info("FastAPI应用实例已创建，可用于uvicorn启动")
    
    return app


def create_app_sync():
    """
    同步包装函数，用于uvicorn工厂模式
    
    解决uvicorn期望同步函数但我们需要异步初始化的问题
    使用nest_asyncio允许在已有事件循环中运行新的事件循环
    
    Returns:
        FastAPI应用实例
    """
    # 使用nest_asyncio运行异步函数
    import asyncio
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    
    app = loop.run_until_complete(create_app())
    return app


async def main():
    """主函数，启动生产环境服务器
    
    使用程序化方式启动uvicorn服务器，而不是命令行方式
    这样可以更好地控制服务器配置和错误处理
    """
    try:
        # 获取服务器配置
        server_config = ServerConfig.from_env()
        
        # 创建应用实例
        app = await create_app()
        
        # 创建并启动服务器
        logger.info(f"正在启动生产环境服务器，端口: {server_config.port}")
        
        # 获取uvicorn配置
        import uvicorn.config
        
        # 创建uvicorn配置
        config = uvicorn.config.Config(
            app=app,
            host="0.0.0.0",
            port=server_config.port,
            log_level="info"
        )
        
        # 创建服务器实例
        uvicorn_server = uvicorn.Server(config)
        
        # 启动服务器
        await uvicorn_server.serve()
        
    except Exception as e:
        logger.error(f"启动服务器时发生错误: {e}", exc_info=True)
        raise

if __name__ == "__main__":
    # 运行主函数
    # 使用asyncio.run启动异步主函数
    asyncio.run(main())
